package org.neu.ccs.secureim;

import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.Socket;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.Vector;
import java.util.concurrent.ConcurrentHashMap;

public class ServerThread implements Runnable {

	Socket sock;
	KeyPair pair;
	static SecureRandom r;
	static KeyFactory kf;
	static ConcurrentHashMap<String, ServerUserInfo> users;
	static Vector<String> loggedIn;
	
	static final int LOGIN = 1;
	static final int REGISTER = 2;
	static final int OKAY = 3;
	static final int ERROR = 4;
	static final int LIST = 5;
	static final int LOCATE = 6;
	static final int LOGOUT = 9;
	
	public ServerThread(Socket sock, KeyPair pair, ConcurrentHashMap<String, ServerUserInfo> users2)
	{
		this.sock = sock;
		this.pair = pair;
		if( r == null )
			r = new SecureRandom();
		if( kf == null )
		{
			try {
				kf = KeyFactory.getInstance("RSA");
			} catch (NoSuchAlgorithmException e) {
				e.printStackTrace();
			}
		}
		if( loggedIn == null )
			loggedIn = new Vector<String>();
		if( users == null )
		{
			users = users2;
		}
	}
	
	@Override
	public void run() {
		try {
			int command = sock.getInputStream().read();
			System.out.println("Command: " + command);
			switch(command)
			{
			case LOGIN:	Login();
						break;
			case LIST: List();
						break;
			case LOCATE: Locate();
						break;
			case LOGOUT: Logout();
			}
		} catch (Exception e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		finally
		{
			((SIMSSecurityManager)System.getSecurityManager()).remove(sock.getInetAddress().getHostAddress());
			try {
				sock.close();
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}
	}
	
	public void Login() throws IOException
	{
		Thread.currentThread().setPriority(Thread.MIN_PRIORITY);
		
		DataInputStream sin = new DataInputStream(sock.getInputStream());
		OutputStream sout = sock.getOutputStream();
		
		int length = NetworkHelper.ReadInt(sin);
		
		if( length > 1500 )
		{
			System.out.println("Unexpected response");
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		
		byte[] received = new byte[length];
		sin.readFully(received);
		
		byte[] decrypted = null;
		try
		{
			decrypted = CryptoHelper.UnpackRSA(pair.getPrivate(), received);
		}
		catch(Exception e)
		{
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		byte[] rand1 = Arrays.copyOfRange(decrypted, decrypted.length - 128, decrypted.length);
		String username = new String(Arrays.copyOfRange(decrypted, 0, decrypted.length - 128));
		if( !username.matches("^[a-zA-Z0-9_]+$") )
		{
			System.out.println("Malformed username " + username + " login dropped");
			return;
		}
		
		System.out.println("Logging in user " + username);
		
		ServerUserInfo inf = users.get(username);
		if( inf == null )
		{
			System.out.println("User does not exist, login dropped");
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		
		byte[] salt = inf.getSalt();
		PublicKey clientKey = inf.getPubKey();
		
		byte[] messageOne = CryptoHelper.SignRSA(pair.getPrivate(), salt, rand1);
		
		byte[] rand2 = new byte[128];
		r.nextBytes(rand2);
		
		ByteArrayOutputStream out = new ByteArrayOutputStream();
		out.write(rand2);
		out.write((int)System.currentTimeMillis());
		out.write(username.getBytes());
		
		byte[] messageTwo = CryptoHelper.SignRSA(pair.getPrivate(), out.toByteArray());
		
		out = new ByteArrayOutputStream();
		out.write(messageOne);
		out.write(messageTwo);
		
		NetworkHelper.WriteInt(sout, out.toByteArray().length);
		sout.write(out.toByteArray());
		
		//384+~400 bytes
		length = NetworkHelper.ReadInt(sin);
		
		if( length > 1500 )
		{
			System.out.println("Unexpected response");
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		
		received = new byte[length];
		sin.readFully(received);
		
		byte[] part1 = Arrays.copyOfRange(received, 0, 388);
		byte[] part2 = Arrays.copyOfRange(received, 388, length);
		
		if( !(CryptoHelper.VerifyRSA(clientKey, part1) && CryptoHelper.VerifyRSA(pair.getPublic(), part2) ) )
		{
			System.out.println("Login failed, incorrect key");
			sout.write(ERROR);
			return;
		}
		
		if( !Arrays.equals(rand2, Arrays.copyOfRange(part1, 260, 388)))
		{
			System.out.println("Challenge failed");
			sout.write(ERROR);
			return;
		}

		inf.setLocation(sock.getInetAddress(), NetworkHelper.BytesToInt(Arrays.copyOfRange(part1, 256, 260)));
		System.out.println("Logged in on port " + inf.getPort());
		if( !loggedIn.contains(username) )
			loggedIn.add(username);
		
		sout.write(OKAY);
	}
	
	public void List() throws IOException
	{
		DataInputStream sin = new DataInputStream(sock.getInputStream());
		OutputStream sout = sock.getOutputStream();
		
		int length = NetworkHelper.ReadInt(sin);
		
		if( length > 1500 )
		{
			System.out.println("Unexpected response");
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		
		byte[] received = new byte[length];
		sin.readFully(received);
		String username = new String(CryptoHelper.UnpackRSA(pair.getPrivate(), received));
		
		ServerUserInfo inf = users.get(username);
		if( inf == null || inf.getLocation() == null )
		{
			System.out.println("User not logged in " + username);
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		
		PublicKey clientKey = inf.getPubKey();
		
		length = NetworkHelper.ReadInt(sin);
		
		if( length > 1500 )
		{
			System.out.println("Unexpected response");
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		
		received = new byte[length];
		sin.readFully(received);
		
		if( !CryptoHelper.VerifyRSA(clientKey, received) )
		{
			System.out.println("Invalid authentication");
			sout.write(ERROR);
			return;
		}
		
		sout.write(OKAY);
		ByteArrayOutputStream out = new ByteArrayOutputStream(); 
		ObjectOutputStream serializer = new ObjectOutputStream(out);
		serializer.writeObject(loggedIn);
		
		byte[] encrypted = CryptoHelper.PackRSA(clientKey, r, out.toByteArray());
		NetworkHelper.WriteInt(sout, encrypted.length);
		sout.write(encrypted);
	}

	public void Locate() throws IOException
	{
		DataInputStream sin = new DataInputStream(sock.getInputStream());
		OutputStream sout = sock.getOutputStream();
		
		int length = NetworkHelper.ReadInt(sin);
		
		if( length > 1500 )
		{
			System.out.println("Unexpected response");
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		
		byte[] received = new byte[length];
		sin.readFully(received);
		String authusername = new String(CryptoHelper.UnpackRSA(pair.getPrivate(), received));
		
		ServerUserInfo inf = users.get(authusername);
		if( inf == null || inf.getLocation() == null )
		{
			System.out.println("User not logged in " + authusername);
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		
		PublicKey clientKey = inf.getPubKey();
		
		length = NetworkHelper.ReadInt(sin);
		
		if( length > 1500 )
		{
			System.out.println("Unexpected response");
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		
		received = new byte[length];
		sin.readFully(received);
		
		String username = new String(Arrays.copyOfRange(received, 256, length));
		
		if( !CryptoHelper.VerifyRSA(clientKey, received) )
		{
			System.out.println("Invalid authentication");
			sout.write(ERROR);
			return;
		}
		
		try
		{
			ServerUserInfo locinf = users.get(username);
			
			if( locinf == null || locinf.getLocation() == null )
			{
				System.out.println("Located user not logged in" + username);
				sout.write(ERROR);
				return;
			}
			
			InetAddress address = locinf.getLocation();
			PublicKey pubKey = locinf.getPubKey();
			
			ByteArrayOutputStream out = new ByteArrayOutputStream(); 
			ObjectOutputStream serializer = new ObjectOutputStream(out);
			serializer.writeObject(address);
			serializer.writeInt(locinf.getPort());
			serializer.writeObject(pubKey);
			sout.write(OKAY);
			byte[] encrypted = CryptoHelper.PackRSA(clientKey, r, out.toByteArray());

			NetworkHelper.WriteInt(sout, encrypted.length);
			sout.write(encrypted);
		}
		catch(Exception e)
		{
			e.printStackTrace();
		}
	}
	
	public void Logout() throws IOException
	{
		System.out.println("logging out");
		DataInputStream sin = new DataInputStream(sock.getInputStream());
		OutputStream sout = sock.getOutputStream();
		
		int length = NetworkHelper.ReadInt(sin);
		
		if( length > 1500 )
		{
			System.out.println("Unexpected response");
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		
		byte[] received = new byte[length];
		sin.readFully(received);
		String authusername = new String(CryptoHelper.UnpackRSA(pair.getPrivate(), received));
		
		ServerUserInfo inf = users.get(authusername);
		if( inf == null || inf.getLocation() == null )
		{
			System.out.println("User not logged in " + authusername);
			NetworkHelper.WriteInt(sout, ERROR);
			return;
		}
		
		PublicKey clientKey = inf.getPubKey();
		
		length = NetworkHelper.ReadInt(sin);
		
		if( length > 1500 )
		{
			System.out.println("Unexpected response");
			NetworkHelper.WriteInt(sout, ERROR);
		}
		
		received = new byte[length];
		sin.readFully(received);
		
		if( CryptoHelper.VerifyRSA(clientKey, received) )
		{
			byte[] timestamp = Arrays.copyOfRange(received, 256, length);
			int istamp = NetworkHelper.BytesToInt(timestamp);
			System.out.println(Math.abs(((int)System.currentTimeMillis()) - istamp));
			if( Math.abs(((int)System.currentTimeMillis()) - istamp) < 5000 )
			{
				users.get(authusername).setLocation(null, 0);
				loggedIn.remove(authusername);
				sout.write(OKAY);
				return;
			}
		}
		else
		{
			System.out.println("Verification failed on logout");
		}
		sout.write(ERROR);
	}
	
}
